import torch
from torch import nn, optim
from torch.utils.data import TensorDataset, DataLoader
import matplotlib.pyplot as plt

# 数据
X = torch.linspace(0, 1, 200).unsqueeze(1)
Y = X ** 6
dataset = TensorDataset(X, Y)
dataloader = DataLoader(dataset, batch_size=20, shuffle=True)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 模型
class DeepModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(1, 128),
            nn.Tanh(),
            nn.Linear(128, 128),
            nn.Tanh(),
            nn.Linear(128, 128),
            nn.Tanh(),
            nn.Linear(128, 1)
        )

    def forward(self, x):
        return self.net(x)



# model = DeepModel().to(device)
#
# print(model)
# exit()

loss_fn = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

epochs = 5000
loss_list = []

for epoch in range(epochs):
    total_loss = 0
    for X_batch, Y_batch in dataloader:
        X_batch = X_batch.to(device)
        Y_batch = Y_batch.to(device)

        y_pred = model(X_batch)
        loss = loss_fn(y_pred, Y_batch)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_loss += loss.item()

    avg_loss = total_loss / len(dataloader)
    loss_list.append(avg_loss)

    if (epoch + 1) % 500 == 0:
        print(f"Epoch {epoch + 1}, Loss={avg_loss:.6f}")

# 可视化
with torch.no_grad():
    pred = model(X.to(device)).cpu()
plt.scatter(X, Y, label='True')
plt.plot(X, pred, color='red', label='Predicted')
plt.legend()
plt.show()
